import numpy as np
import pickle
from Config.Config import ACTNET200V13_PKL
from Tools.utils import nms, recall_vs_iou_thresholds, get_gt, convert

# RANKED_PROPOSALS_FILE = '/mnt/md1/Experiments/PYTSN_Test2/VAL_BACKUP/ranked_proposals.pkl'
RANKED_PROPOSALS_FILE = '/mnt/md1/Experiments/PYTSN_Test7/ranked_proposals.pkl'

print('ranked_file: {}'.format(RANKED_PROPOSALS_FILE))

with open(RANKED_PROPOSALS_FILE,'rb') as f:
    ranked_proposals = pickle.load(f)

with open(ACTNET200V13_PKL, 'rb') as f:
    gt = pickle.load(f)['database']

proposal_at_1 = {'s-init':[],'s-end':[],'score':[],'label':[],'video-id':[]}

for vid, proposal in ranked_proposals.items():

    proposal = np.asarray(proposal)
    # 处于未知的原因需要特别判断一下proposal的形状
    if len(proposal.shape) == 3: proposal = proposal.reshape((proposal.shape[0]*proposal.shape[1], proposal.shape[2]))
    else: proposal = np.concatenate(proposal)

    keep_ind = nms(proposal, proposal[:,2], 0.45)
    proposal = proposal[keep_ind,:]

    proposal_at_1['s-init'].append(proposal[0,0])
    proposal_at_1['s-end'].append(proposal[0,1])
    proposal_at_1['score'].append(proposal[0,2])
    proposal_at_1['video-id'].append('v_'+vid)

gt2 = get_gt(gt)
iou_thrs = np.arange(0.1, 1.0, 0.1)
recall_at_1 = recall_vs_iou_thresholds(convert(proposal_at_1), gt2, iou_threshold=iou_thrs)

print(np.array_str(np.vstack([iou_thrs, recall_at_1]), precision=4, suppress_small=True),flush=True)
